import pgzrun
import random
import time

ROWS = 50
COLS = 50
CELL_SIZE = 15

HEIGHT = (ROWS * CELL_SIZE)
WIDTH  = (COLS * CELL_SIZE)

BACK_COLOR = [0, 0, 127]

def Apply(w, func):
    for r in range(ROWS):
        for c in range(COLS):
            w[r][c] = func(r, c)

def GenCellColor():
    return [random.randint(0, 200), random.randint(0, 200), random.randint(0, 200)]

def init_map(rows, cols):
    return [[[0, BACK_COLOR] for c in range(cols)] for r in range(rows)]

world = init_map(ROWS, COLS)
next_world = init_map(ROWS, COLS)

def random_map(w):
    Apply(w, lambda r, c: [int(random.randint(0, 7) == 0), GenCellColor()])

def draw_cell(w, r, c):
    cx = CELL_SIZE * c
    cy = CELL_SIZE * r
    cell_rect = Rect((cx, cy), (CELL_SIZE, CELL_SIZE))
    screen.draw.filled_rect(cell_rect, tuple(w[r][c][1]))
    return w[r][c]

def count_neighbors(w, r, c):
    sum = 0
    clr = [0, 0, 0]
    for nr in range(r - 1, r + 1 + 1):
        for nc in range(c - 1, c + 1 + 1):
            if nr == r and nc == c:
                continue

            if nr < 0 or nr >= ROWS:
                continue

            if nc < 0 or nc >= COLS:
                continue

            if w[nr][nc][0]:
                sum += 1
                clr[0] += w[nr][nc][1][0]
                clr[1] += w[nr][nc][1][1]
                clr[2] += w[nr][nc][1][2]
    if sum == 0:
        return sum, clr
    else:
        return sum, [clr[0] // sum, clr[1] // sum, clr[2] // sum]

def next_cell(w, r, c):
    n, clr = count_neighbors(w, r, c)

    if n == 3:
        return [1, clr]
    elif n == 2:
        return w[r][c]
    else:
        return [0, BACK_COLOR]

def update():
    Apply(next_world, lambda r, c: next_cell(world, r, c))
    Apply(world, lambda r, c: next_world[r][c])
             
    time.sleep(0.3)

def draw():
    screen.fill(tuple(BACK_COLOR))
    Apply(world, lambda r, c: (draw_cell(world, r, c) if world[r][c][0] else [0, BACK_COLOR]))

random_map(world)

pgzrun.go()